import os
from unittest.mock import patch

from library_test.base_test import BaseTest


class AnsibleModuleMocker:
    def __init__(self, params: dict):
        self.params = params

    def get_bin_path(self, path):
        return True


class BaseTestNpuChecks(BaseTest):

    TESTCASE_DIR = os.path.join(os.path.dirname(__file__), "testcase")

    @classmethod
    def setUpClass(cls) -> None:
        super().setUpClass()
        cls.mock_func()

        run_cmd_patcher = patch(cls.get_module_path() + ".util.run_cmd")
        cls.cmd_mocker = run_cmd_patcher.start()

        record_error_patcher = patch(cls.get_module_path() + ".util.record_error")
        cls.record_mocker = record_error_patcher.start()

        module = AnsibleModuleMocker({})
        cls.error_msg = []
        from ascend_deployer.module_utils.check_library_utils.npu_checks import NPUCheck
        cls.npu_check = NPUCheck(module, cls.error_msg)

    @classmethod
    def mock_func(cls):
        mock_get_card = patch(cls.get_module_path() + ".util.get_card")
        mock_get_card.start()
        mock_get_card.return_value = "910_93"

        mock_get_npu_info = patch(cls.get_module_path() + ".common_info.get_npu_info")
        mock_get_npu_info.start()
        mock_get_npu_info.return_value = {
            "card": "A300i-pro",
            "model": "A300i-pro",
            "scene": "a300i",
            "product": "A300i",
        }

    @classmethod
    def get_module_path(cls):
        return "ascend_deployer.module_utils.check_library_utils.npu_checks"

    
    def clear(self):
        self.error_msg = []

    def test_is_occupied_by_process(self):
        self.cmd_mocker.return_value = ""
        self.assertFalse(self.npu_check.is_occupied_by_process())

        self.cmd_mocker.return_value = b'''
            +--------------------------------------------------------------------------------------------------------+
            | npu-smi 25.3.rc1.b020                            Version: 25.3.rc1.b020                                |
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Name                  | Health          | Power(W)     Temp(C)           Hugepages-Usage(page) |
            | Chip    Device                | Bus-Id          | AICore(%)    Memory-Usage(MB)                        |
            +===============================+=================+======================================================+
            | 3       310P3                 | OK           | NA           56                0     / 0             |
            | 0       0                     | 0000:03:00.0    | 0            1883 / 23047                            |
            +===============================+=================+======================================================+
            | 5       310P3                 | OK              | NA           57                0     / 0             |
            | 0       1                     | 0000:81:00.0    | 0            1881 / 23047                            |
            +===============================+=================+======================================================+
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Chip                  | Process id      | Process name             | Process memory(MB)        |
            +===============================+=================+======================================================+
            | 3       0                     | 3860114         | mindie_llm_back          | 58427                     |
            +===============================+=================+======================================================+
            | No running processes found in NPU 5                                                                    |
            +===============================+=================+======================================================+
        '''
        self.assertTrue(self.npu_check.is_occupied_by_process())

    def test_is_occupied_by_docker(self):
        # no out
        self.cmd_mocker.return_value = ""
        self.assertFalse(self.npu_check.is_occupied_by_docker("file"))

        # decode failed
        self.cmd_mocker.return_value = "decode failed"
        with self.assertRaises(RuntimeError) as context:
            self.npu_check.is_occupied_by_docker("file")
        self.assertIn("Get docker container ids failed", str(context.exception))

        # no container id
        self.cmd_mocker.return_value = b""
        self.assertFalse(self.npu_check.is_occupied_by_docker("file"))
        
        # docker inspect no out
        self.cmd_mocker.return_value = b"22454f5671bb\n89dd89ad34dd"
        self.assertTrue(self.npu_check.is_occupied_by_docker("file"))

    def test_check_npu_health_with_no_command(self):
        self.clear()
        self.cmd_mocker.return_value = b''''''
        self.npu_check.check_npu_health()
        self.assertEqual(len(self.error_msg), 0)

    def test_check_npu_health_with_success(self):
        self.clear()
        self.cmd_mocker.return_value = b'''
            +--------------------------------------------------------------------------------------------------------+
            | npu-smi 25.3.rc1.b020                            Version: 25.3.rc1.b020                                |
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Name                  | Health          | Power(W)     Temp(C)           Hugepages-Usage(page) |
            | Chip    Device                | Bus-Id          | AICore(%)    Memory-Usage(MB)                        |
            +===============================+=================+======================================================+
            | 3       310P3                 | OK              | NA           56                0     / 0             |
            | 0       0                     | 0000:03:00.0    | 0            1883 / 23047                            |
            +===============================+=================+======================================================+
            | 5       310P3                 | OK              | NA           57                0     / 0             |
            | 0       1                     | 0000:81:00.0    | 0            1881 / 23047                            |
            +===============================+=================+======================================================+
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Chip                  | Process id      | Process name             | Process memory(MB)        |
            +===============================+=================+======================================================+
            | No running processes found in NPU 3                                                                    |
            +===============================+=================+======================================================+
            | No running processes found in NPU 5                                                                    |
            +===============================+=================+======================================================+
        '''
        self.npu_check.check_npu_health()
        self.assertEqual(len(self.error_msg), 0)

    def test_check_npu_health_with_910_value(self):
        self.clear()
        self.cmd_mocker.return_value = b'''
            +--------------------------------------------------------------------------------------------------------+
            | npu-smi 25.3.rc1.b020                            Version: 25.3.rc1.b020                                |
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Name                  | Health          | Power(W)     Temp(C)           Hugepages-Usage(page) |
            | Chip    Device                | Bus-Id          | AICore(%)    Memory-Usage(MB)                        |
            +===============================+=================+======================================================+
            | 3       310P3                 | OK              | NA           56                0     / 0             |
            | 0       0                     | 0000:03:00.0    | 0            1910 / 23047                            |
            +===============================+=================+======================================================+
            | 5       310P3                 | OK              | NA           57                0     / 0             |
            | 0       1                     | 0000:81:00.0    | 0            1881 / 23047                            |
            +===============================+=================+======================================================+
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Chip                  | Process id      | Process name             | Process memory(MB)        |
            +===============================+=================+======================================================+
            | No running processes found in NPU 3                                                                    |
            +===============================+=================+======================================================+
            | No running processes found in NPU 5                                                                    |
            +===============================+=================+======================================================+
        '''
        self.npu_check.check_npu_health()
        self.assertEqual(len(self.error_msg), 0)

    def test_check_npu_health_with_error(self):
        self.cmd_mocker.return_value = b'''
            +--------------------------------------------------------------------------------------------------------+
            | npu-smi 25.3.rc1.b020                            Version: 25.3.rc1.b020                                |
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Name                  | Health          | Power(W)     Temp(C)           Hugepages-Usage(page) |
            | Chip    Device                | Bus-Id          | AICore(%)    Memory-Usage(MB)                        |
            +===============================+=================+======================================================+
            | 3       310P3                 | Error           | NA           56                0     / 0             |
            | 0       0                     | 0000:03:00.0    | 0            1883 / 23047                            |
            +===============================+=================+======================================================+
            | 5       310P3                 | OK              | NA           57                0     / 0             |
            | 0       1                     | 0000:81:00.0    | 0            1881 / 23047                            |
            +===============================+=================+======================================================+
            +-------------------------------+-----------------+------------------------------------------------------+
            | NPU     Chip                  | Process id      | Process name             | Process memory(MB)        |
            +===============================+=================+======================================================+
            | No running processes found in NPU 3                                                                    |
            +===============================+=================+======================================================+
            | No running processes found in NPU 5                                                                    |
            +===============================+=================+======================================================+
        '''
        self.npu_check.check_npu_health()
        self.record_mocker.assert_called()
